#include "tracker.h"

#include <Eigen/Core>
#include <Eigen/Dense>

void OpticalFlowTracker::calculateOpticalFlow(const Range &range)
{
    // parameters
    int half_patch_size = 4;
    int iterations = 10;
    for (size_t i = range.start; i < range.end; i++)
    {
        auto kp = kp1[i];
        double dx = 0, dy = 0; // dx,dy need to be estimated
        if (has_initial)
        {
            dx = kp2[i].pt.x - kp.pt.x;
            dy = kp2[i].pt.y - kp.pt.y;
        }

        double cost = 0, lastCost = 0;
        bool succ = true; // indicate if this point succeeded

        // Gauss-Newton iterations
        //  
        // (Δx*,Δy*) = argmin 1/2||I1(x,y)-I2(x+Δx, y+Δy)||^2
        //             Δx,Δy
        // 
        // e(Δx,Δy)  = I1(x,y)-I2(x+Δx, y+Δy)
        // 
        //               [∂I2 ∂I2]|
        //        J' = - |--- ---||
        //               [∂Δx ∂Δy]|(x+Δx, y+Δy)
        Eigen::Matrix2d H = Eigen::Matrix2d::Zero(); // hessian: JJ'
        Eigen::Vector2d g = Eigen::Vector2d::Zero(); // e*J
        Eigen::Vector2d J;                           // jacobian
        for (int iter = 0; iter < iterations; iter++)
        {
            if (inverse == false)
            {
                H = Eigen::Matrix2d::Zero();
                g = Eigen::Vector2d::Zero();
            }
            else
            {
                // only reset
                g = Eigen::Vector2d::Zero();
            }

            cost = 0;

            // compute cost and jacobian
            for (int x = -half_patch_size; x < half_patch_size; x++)
                for (int y = -half_patch_size; y < half_patch_size; y++)
                {
                    double error = GetGrayscale(img1, kp.pt.x + x, kp.pt.y + y) -
                                   GetGrayscale(img2, kp.pt.x + x + dx, kp.pt.y + y + dy);
                    // Jacobian
                    if (inverse == false)
                    {
                        J = -1.0 * Eigen::Vector2d(
                                       0.5 * (GetGrayscale(img2, kp.pt.x + dx + x + 1, kp.pt.y + dy + y) -
                                              GetGrayscale(img2, kp.pt.x + dx + x - 1, kp.pt.y + dy + y)),
                                       0.5 * (GetGrayscale(img2, kp.pt.x + dx + x, kp.pt.y + dy + y + 1) -
                                              GetGrayscale(img2, kp.pt.x + dx + x, kp.pt.y + dy + y - 1)));
                    }
                    else if (iter == 0)
                    {
                        // in inverse mode, J keeps same for all iterations
                        // NOTE this J does not change when dx, dy is updated, so we can store it and only compute error
                        J = -1.0 * Eigen::Vector2d(
                                       0.5 * (GetGrayscale(img1, kp.pt.x + x + 1, kp.pt.y + y) -
                                              GetGrayscale(img1, kp.pt.x + x - 1, kp.pt.y + y)),
                                       0.5 * (GetGrayscale(img1, kp.pt.x + x, kp.pt.y + y + 1) -
                                              GetGrayscale(img1, kp.pt.x + x, kp.pt.y + y - 1)));
                    }
                    // compute H, g and set cost;
                    g += -error * J;
                    cost += error * error;
                    if (inverse == false || iter == 0)
                    {
                        // also update H
                        H += J * J.transpose();
                    }
                }

            // compute update
            Eigen::Vector2d update = H.ldlt().solve(g);

            if (std::isnan(update[0]))
            {
                // sometimes occurred when we have a black or white patch and H is irreversible
                cout << "update is nan" << endl;
                succ = false;
                break;
            }

            if (iter > 0 && cost > lastCost)
            {
                break;
            }

            // update dx, dy
            dx += update[0];
            dy += update[1];
            lastCost = cost;
            succ = true;

            if (update.norm() < 1e-2)
            {
                // converge
                break;
            }
        }

        success[i] = succ;

        // set kp2
        kp2[i].pt = kp.pt + Point2f(dx, dy);
    }
}
